#include "style.h"

void drawSatur(const char * input, const char *cname, const Bool_t kecor=0, const Bool_t kNoElectron = kTRUE)
{
  //
  // 1. AliMathBase::BetheBlochAleph is the input fuction for the mean number of collision per cm in the TPC simulation
  //    (user defined)
  //    for primary collision above 10 KeV secondaries are produced by Geant 
  //    Question how is the <dEdx>/<N> is as a function of input <N>
  //       As a model we are using the scaling function
  //    What we can see there is the convolution of several effects but dominant is
  //         the truncated mean (but not only also detector effects are present)
  //

  //================================>
  //important variables
  const Double_t momElossCorrection = 4.6; 
  //from bb solid - should be cross-checked - do we need to make fit? form data extracted vaule 4.1 
  //              - eta  range  used should be part of the tree  - to have correct ELoss correction
  //<================================
  // TFile f("high.root", "update");
  TFile f(input, "update");
  
  if(!tree->GetBranch("meanbgproton")){
    printf("no meanbgproton\n");
    tree->Print();
    exit(1);
  }

  const Int_t kcolors[]={kRed, kBlue, kBlack, kMagenta};
  const Int_t sty[]={20,29,22,21};
  const char * pname[]={"proton","pion","kaon","electron"};
  const Double_t mass[]={0.938272, 0.139570, 0.493677, 5.109989e-04 };
        
  //temporary energy loss correction, particle type dependent
  //  const Double_t tmpcor[]={5.89975e-03, 7.92038e-03, 7.42592e-03, 7.92038e-03};
  //const Double_t tmpcor[]={8.89975e-03, 7.92038e-03, 7.42592e-03, 7.92038e-03};
  TString var, icut;
  Int_t entries = -999;

  TCanvas * working= new TCanvas;

  Int_t ntype = 4;
  TGraphErrors *grBB[4];
  for(Int_t ii=0; ii<ntype; ii++){
    if(!kecor){
      tree->SetAlias(Form("bg%s",pname[ii]), Form("meanbg%s", pname[ii]));
    } 
    else{
      //only here use bbsolid for mom energy loss correction
      var = Form("sqrt(pow(sqrt(meanp*meanp+%f*%f)-%f*AliExternalTrackParam::BetheBlochSolid(meanp/%f),2)-pow(%f,2))/%f", mass[ii], mass[ii], /*tmpcor[ii]*/ momElossCorrection, mass[ii], mass[ii], mass[ii]);
      printf("check1 %d: %s\n", ii, var.Data());
      tree->SetAlias(Form("bg%s", pname[ii]), var);
    }
    
    var=Form("0.05+e%smean/AliMathBase::BetheBlochAleph(bg%s):p%smean/AliMathBase::BetheBlochAleph(bg%s):1/AliMathBase::BetheBlochAleph(bg%s)", pname[ii], pname[ii], pname[ii], pname[ii], pname[ii]);
    printf("check2 %d: %s\n", ii, var.Data());
    //tree->SetScanField(0);

    icut = Form("integral%s>20 & meanp>0.5&&bg%s>1", pname[ii], pname[ii]);
    printf("check2cut %d: %s\n", ii, icut.Data());
    //tree->Scan(Form("integral%s:meanp:bg%s", pname[ii], pname[ii]),Form("integral%s>20 & meanp>0.5&&bg%s>1", pname[ii], pname[ii]),"GOFF"); 

    //important to check the cut here!! 
    //const Int_t entries=tree->Draw(var,Form("integral%s>20 & meanp>0.5&&bg%s>1", pname[ii], pname[ii]),"GOFF"); 
    entries=tree->Draw(var, icut,"GOFF"); 
    grBB[ii] = new TGraphErrors(entries, tree->GetV3(), tree->GetV2(), 0, tree->GetV1());
  }
  
  TF1 fscaleBB("fscaleBB","[0]+[1]*(1-x)+[2]*(1-x)**2",0,1);

  grBB[0]->GetXaxis()->SetLimits(0.58,1.05);
  Double_t p0=0,p1=0,p2=0;
  
  {for (Int_t i=0; i<ntype; i++){
      grBB[i]->SetMarkerStyle(sty[i]);
      grBB[i]->SetMarkerColor(kcolors[i]);
      grBB[i]->SetLineColor(kcolors[i]);
      fscaleBB.SetLineColor(kcolors[i]);
      grBB[i]->GetXaxis()->SetTitle("1/BB_{input}(1/MIP)");
      grBB[i]->GetYaxis()->SetTitle("dEdx_{m}/BB_{input}");
      grBB[i]->Draw(i?"lp":"alp");
      //grBB[i]->Fit("fscaleBB","W","W");
      /*
      p0+=fscaleBB.GetParameter(0)/ntype;
      p1+=fscaleBB.GetParameter(1)/ntype;
      p2+=fscaleBB.GetParameter(2)/ntype;
      */
    }
  }

  TCanvas *ctot=new TCanvas;
  TGraphErrors *grtot=new TGraphErrors;

  Double_t p0=0,p1=0,p2=0;
  
  for(Int_t itype =0; itype<ntype; itype++){
    for(Int_t ip = 0; ip<grBB[itype]->GetN(); ip++){
      const Int_t ntot = grtot->GetN();
      grtot->SetPoint(     ntot, grBB[itype]->GetX()[ip],  grBB[itype]->GetY()[ip]);
      grtot->SetPointError(ntot, grBB[itype]->GetEX()[ip], grBB[itype]->GetEY()[ip]);
    }
  }
  grtot->Draw("alp");
  grtot->Fit("fscaleBB");//,"W","W");
  p0=fscaleBB.GetParameter(0);
  p1=fscaleBB.GetParameter(1);
  p2=fscaleBB.GetParameter(2);
  //
  // 2. Create an modified BB, check how well is the data described by modified function
  //  
  TGraphErrors *grBB2[4];
  TGraphErrors *grBB3[4];
  for(Int_t ii=0; ii<ntype; ii++){
    var=Form("AliMathBase::BetheBlochAleph(bg%s)*(%f+%f*(1-1/AliMathBase::BetheBlochAleph(bg%s))+ %f*(1-1/AliMathBase::BetheBlochAleph(bg%s))**2 )", pname[ii], p0, p1, pname[ii], p2, pname[ii]);
    printf("check3 %d: %s\n", ii, var.Data());
    tree->SetAlias(Form("BB%s",pname[ii]),var);

    var=Form("e%smean/BB%s:p%smean/BB%s:bg%s", pname[ii], pname[ii], pname[ii], pname[ii], pname[ii]);
    printf("check4 %d: %s\n", ii, var.Data());
    entries = tree->Draw(var,Form("bg%s>0.5",pname[ii]),"goff");
    grBB2[ii]= new TGraphErrors(entries, tree->GetV3(), tree->GetV2(), 0, tree->GetV1());

    var=Form("e%smean/BB%s:p%smean/BB%s:meanp", pname[ii], pname[ii], pname[ii], pname[ii], pname[ii]);
    printf("check5 %d: %s\n", ii, var.Data());
    entries = tree->Draw(var,"","goff");
    grBB3[ii] = new TGraphErrors(entries, tree->GetV3(), tree->GetV2(), 0, tree->GetV1());
    grBB3[ii]->SetName(Form("gr3%s",pname[ii]));
  }

  //get the 1/p correction
  TF1 fpt("fpt","[0]+[1]/x",0.1,4);
  fpt.SetParameters(1,0);
  
  Double_t a0=-999, a1=-999;

  TCanvas *ccomb=new TCanvas;
  gPad->SetLogx();
  TGraphErrors *grcomb=new TGraphErrors;

  for(Int_t itype =0; itype<ntype; itype++){
    for(Int_t ip = 0; ip<grBB3[itype]->GetN(); ip++){
      const Int_t ncomb = grcomb->GetN();
      grcomb->SetPoint(     ncomb, grBB3[itype]->GetX()[ip],  grBB3[itype]->GetY()[ip]);
      grcomb->SetPointError(ncomb, grBB3[itype]->GetEX()[ip], grBB3[itype]->GetEY()[ip]);
    }
  }
  grcomb->Draw("ap");
  grcomb->Fit("fpt","","",0.4,20);
  a0=fpt.GetParameter(0);
  a1=fpt.GetParameter(1);

  TGraphErrors *grBB4[4];
  TGraphErrors *grBB5[4];
  for(Int_t ii=0; ii<ntype; ii++){
    var=Form("AliMathBase::BetheBlochAleph(bg%s)*(%f+%f*(1-1/AliMathBase::BetheBlochAleph(bg%s))+ %f*(1-1/AliMathBase::BetheBlochAleph(bg%s))**2 )*(%f+%f/meanp)", pname[ii], p0, p1, pname[ii], p2, pname[ii], a0, a1);
    printf("check6 %d: %s\n", ii, var.Data());
    tree->SetAlias(Form("BB4%s",pname[ii]),var);

    var=Form("e%smean/BB4%s:p%smean/BB4%s:meanp", pname[ii], pname[ii], pname[ii], pname[ii], pname[ii]);
    printf("check7 %d: %s\n", ii, var.Data());
    entries = tree->Draw(var,"","goff");
    grBB4[ii] = new TGraphErrors(entries, tree->GetV3(), tree->GetV2(), 0, tree->GetV1());
    grBB4[ii]->SetName(Form("gr4%s",pname[ii]));

    var=Form("e%smean/BB4%s:p%smean/BB4%s:bg%s", pname[ii], pname[ii], pname[ii], pname[ii], pname[ii]);
    printf("check8 %d: %s\n", ii, var.Data());
    entries = tree->Draw(var,Form("bg%s>0.5",pname[ii]),"goff");
    grBB5[ii]= new TGraphErrors(entries, tree->GetV3(), tree->GetV2(), 0, tree->GetV1());
  }
  //

  //drawing all

  grBB[0]->SetMaximum(55.);
  grBB[0]->SetMinimum(40);
  //grBB2[0]->GetXaxis()->SetRangeUser(0.5,50);

  grBB2[0]->SetMaximum(1.02);
  grBB2[0]->SetMinimum(0.98);
  grBB2[0]->GetXaxis()->SetLimits(0.2,50000);

  grBB3[0]->SetMaximum(1.02);
  grBB3[0]->SetMinimum(0.98);
  grBB3[0]->GetXaxis()->SetLimits(0.2,50);

  grBB4[0]->SetMaximum(1.02);
  grBB4[0]->SetMinimum(0.98);
  grBB4[0]->GetXaxis()->SetLimits(0.2,50);
  
  grBB5[0]->SetMaximum(1.02);
  grBB5[0]->SetMinimum(0.98);
  grBB5[0]->GetXaxis()->SetLimits(0.2,50000);

  TCanvas  * canvas= new TCanvas(cname, cname,1000,800);
  canvas->Divide(2,3);
  
  TLegend * lg1 = new TLegend(0.8, 0.1, 0.9, 0.5);
  ResetStyle(lg1);
  TLegend * lg2 = new TLegend(0.8, 0.1, 0.9, 0.5);
  ResetStyle(lg2);
  TLegend * lg3 = new TLegend(0.8, 0.1, 0.9, 0.5);
  ResetStyle(lg3);
  TLegend * lg4 = new TLegend(0.8, 0.1, 0.9, 0.5);
  ResetStyle(lg4);
  TLegend * lg5 = new TLegend(0.8, 0.1, 0.9, 0.5);
  ResetStyle(lg5);

  //electron not drawn due to low statistics
  const Int_t ndrawtype = kNoElectron? 3 : ntype; 
  {for (Int_t i=0; i<ndrawtype; i++){
      canvas->cd(1)->SetLogx(0);
      gPad->SetGrid();
      grBB[i]->Draw(i?"lp":"alp");
      lg1->AddEntry(grBB[i], pname[i],"lp");
      lg1->Draw();
      
      canvas->cd(2+1)->SetLogx(1);
      gPad->SetGrid();
      grBB3[i]->SetMarkerStyle(sty[i]);
      grBB3[i]->SetMarkerColor(kcolors[i]);
      grBB3[i]->SetLineColor(kcolors[i]);
      grBB3[i]->GetXaxis()->SetTitle("p (GeV/c)");
      grBB3[i]->GetYaxis()->SetTitle("dEdx_{m}/(BB_{input}*c_{satur})");
      grBB3[i]->Draw(i?"lp":"alp");
      //grBB3[i]->Fit("fpt","","",0.4,20);
      lg3->AddEntry(grBB3[i], pname[i],"lp");
      lg3->Draw();

      canvas->cd(3+1)->SetLogx(1);
      gPad->SetGrid();
      grBB4[i]->SetMarkerStyle(sty[i]);
      grBB4[i]->SetMarkerColor(kcolors[i]);
      grBB4[i]->SetLineColor(kcolors[i]);
      grBB4[i]->GetXaxis()->SetTitle("p (GeV/c)");
      grBB4[i]->GetYaxis()->SetTitle("dEdx_{m}/(BB_{input}*c_{satur}*c_{1/p})");
      grBB4[i]->Draw(i?"lp":"alp");
      lg4->AddEntry(grBB4[i], pname[i],"lp");
      lg4->Draw();

      canvas->cd(4+1)->SetLogx(1);
      gPad->SetGrid();
      
      grBB2[i]->SetMarkerStyle(sty[i]);
      grBB2[i]->SetMarkerColor(kcolors[i]);
      grBB2[i]->SetLineColor(kcolors[i]);
      grBB2[i]->GetXaxis()->SetTitle("#beta#gamma");
      grBB2[i]->GetYaxis()->SetTitle("dEdx_{m}/(BB_{input}*c_{satur})");
      grBB2[i]->Draw(i?"lp":"alp");
      lg2->AddEntry(grBB2[i], pname[i],"lp");
      lg2->Draw();

      canvas->cd(5+1)->SetLogx(1);
      gPad->SetGrid();
      
      grBB5[i]->SetMarkerStyle(sty[i]);
      grBB5[i]->SetMarkerColor(kcolors[i]);
      grBB5[i]->SetLineColor(kcolors[i]);
      grBB5[i]->GetXaxis()->SetTitle("#beta#gamma");
      grBB5[i]->GetYaxis()->SetTitle("dEdx_{m}/(BB_{input}*c_{satur}*c_{1/p})");
      grBB5[i]->Draw(i?"lp":"alp");
      lg5->AddEntry(grBB5[i], pname[i],"lp");
      lg5->Draw();
    }
  }


  canvas->Write();
}


